{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "<h1> Create Keras Wide-and-Deep model </h1>\n",
    "\n",
    "This notebook illustrates:\n",
    "<ol>\n",
    "<li> Creating a model using Keras. This requires TensorFlow 2.1\n",
    "</ol>"
   ]
  },
  {

"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Ensure the right version of Tensorflow is installed.\n",
"!pip freeze | grep tensorflow==2.1"
]
},
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# change these to try this notebook out\n",
    "BUCKET = 'cloud-training-demos-ml'\n",
    "PROJECT = 'cloud-training-demos'\n",
    "REGION = 'us-central1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['BUCKET'] = BUCKET\n",
    "os.environ['PROJECT'] = PROJECT\n",
    "os.environ['REGION'] = REGION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if ! gcloud storage ls | grep -q gs://${BUCKET}/; then\n",
    "  gcloud storage buckets create --location=${REGION} gs://${BUCKET}\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "ls *.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Create Keras model\n",
    "<p>\n",
    "First, write an input_fn to read the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import shutil\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# Determine CSV, label, and key columns\n",
    "CSV_COLUMNS = 'weight_pounds,is_male,mother_age,plurality,gestation_weeks,key'.split(',')\n",
    "LABEL_COLUMN = 'weight_pounds'\n",
    "KEY_COLUMN = 'key'\n",
    "\n",
    "# Set default values for each CSV column. Treat is_male and plurality as strings.\n",
    "DEFAULTS = [[0.0], ['null'], [0.0], ['null'], [0.0], ['nokey']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def features_and_labels(row_data):\n",
    "    for unwanted_col in ['key']:\n",
    "        row_data.pop(unwanted_col)\n",
    "    label = row_data.pop(LABEL_COLUMN)\n",
    "    return row_data, label  # features, label\n",
    "\n",
    "# load the training data\n",
    "def load_dataset(pattern, batch_size=1, mode=tf.estimator.ModeKeys.EVAL):\n",
    "  dataset = (tf.data.experimental.make_csv_dataset(pattern, batch_size, CSV_COLUMNS, DEFAULTS)\n",
    "             .map(features_and_labels) # features, label\n",
    "             )\n",
    "  if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        dataset = dataset.shuffle(1000).repeat()\n",
    "  dataset = dataset.prefetch(1) # take advantage of multi-threading; 1=AUTOTUNE\n",
    "  return dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Next, define the feature columns. mother_age and gestation_weeks should be numeric.\n",
    "The others (is_male, plurality) should be categorical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Build a Keras wide-and-deep model using its Functional API\n",
    "def rmse(y_true, y_pred):\n",
    "    return tf.sqrt(tf.reduce_mean(tf.square(y_pred - y_true))) \n",
    "\n",
    "# Helper function to handle categorical columns\n",
    "def categorical_fc(name, values):\n",
    "    orig = tf.feature_column.categorical_column_with_vocabulary_list(name, values)\n",
    "    wrapped = tf.feature_column.indicator_column(orig)\n",
    "    return orig, wrapped\n",
    "\n",
    "def build_wd_model(dnn_hidden_units = [64, 32], nembeds = 3):\n",
    "    # input layer\n",
    "    deep_inputs = {\n",
    "        colname : tf.keras.layers.Input(name=colname, shape=(), dtype='float32')\n",
    "           for colname in ['mother_age', 'gestation_weeks']\n",
    "    }\n",
    "    wide_inputs = {\n",
    "        colname : tf.keras.layers.Input(name=colname, shape=(), dtype='string')\n",
    "           for colname in ['is_male', 'plurality']        \n",
    "    }\n",
    "    inputs = {**wide_inputs, **deep_inputs}\n",
    "    \n",
    "    # feature columns from inputs\n",
    "    deep_fc = {\n",
    "        colname : tf.feature_column.numeric_column(colname)\n",
    "           for colname in ['mother_age', 'gestation_weeks']\n",
    "    }\n",
    "    wide_fc = {}\n",
    "    is_male, wide_fc['is_male'] = categorical_fc('is_male', ['True', 'False', 'Unknown'])\n",
    "    plurality, wide_fc['plurality'] = categorical_fc('plurality',\n",
    "                      ['Single(1)', 'Twins(2)', 'Triplets(3)',\n",
    "                       'Quadruplets(4)', 'Quintuplets(5)','Multiple(2+)'])\n",
    "    \n",
    "    # bucketize the float fields. This makes them wide\n",
    "    age_buckets = tf.feature_column.bucketized_column(deep_fc['mother_age'],\n",
    "                                                     boundaries=np.arange(15,45,1).tolist())\n",
    "    wide_fc['age_buckets'] = tf.feature_column.indicator_column(age_buckets)\n",
    "    gestation_buckets = tf.feature_column.bucketized_column(deep_fc['gestation_weeks'],\n",
    "                                                     boundaries=np.arange(17,47,1).tolist())\n",
    "    wide_fc['gestation_buckets'] = tf.feature_column.indicator_column(gestation_buckets)\n",
    "    \n",
    "    # cross all the wide columns. We have to do the crossing before we one-hot encode\n",
    "    crossed = tf.feature_column.crossed_column(\n",
    "        [is_male, plurality, age_buckets, gestation_buckets], hash_bucket_size=20000)\n",
    "    deep_fc['crossed_embeds'] = tf.feature_column.embedding_column(crossed, nembeds)\n",
    "\n",
    "    # the constructor for DenseFeatures takes a list of numeric columns\n",
    "    # The Functional API in Keras requires that you specify: LayerConstructor()(inputs)\n",
    "    wide_inputs = tf.keras.layers.DenseFeatures(wide_fc.values(), name='wide_inputs')(inputs)\n",
    "    deep_inputs = tf.keras.layers.DenseFeatures(deep_fc.values(), name='deep_inputs')(inputs)\n",
    "        \n",
    "    # hidden layers for the deep side\n",
    "    layers = [int(x) for x in dnn_hidden_units]\n",
    "    deep = deep_inputs\n",
    "    for layerno, numnodes in enumerate(layers):\n",
    "        deep = tf.keras.layers.Dense(numnodes, activation='relu', name='dnn_{}'.format(layerno+1))(deep)        \n",
    "    deep_out = deep\n",
    "    \n",
    "    # linear model for the wide side\n",
    "    wide_out = tf.keras.layers.Dense(10, activation='relu', name='linear')(wide_inputs)\n",
    "   \n",
    "    # concatenate the two sides\n",
    "    both = tf.keras.layers.concatenate([deep_out, wide_out], name='both')\n",
    "\n",
    "    # final output is a linear activation because this is regression\n",
    "    output = tf.keras.layers.Dense(1, activation='linear', name='weight')(both)\n",
    "    model = tf.keras.models.Model(inputs, output)\n",
    "    model.compile(optimizer='adam', loss='mse', metrics=[rmse, 'mse'])\n",
    "    return model\n",
    "\n",
    "print(\"Here is our Wide-and-Deep architecture so far:\\n\")\n",
    "model = build_wd_model()\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the DNN using the Keras plot_model utility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.utils.plot_model(model, 'wd_model.png', show_shapes=False, rankdir='LR')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Train and evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "TRAIN_BATCH_SIZE = 32\n",
    "NUM_TRAIN_EXAMPLES = 10000 * 5 # training dataset repeats, so it will wrap around\n",
    "NUM_EVALS = 5  # how many times to evaluate\n",
    "NUM_EVAL_EXAMPLES = 10000 # enough to get a reasonable sample, but not so much that it slows down\n",
    "\n",
    "trainds = load_dataset('train*', TRAIN_BATCH_SIZE, tf.estimator.ModeKeys.TRAIN)\n",
    "evalds = load_dataset('eval*', 1000, tf.estimator.ModeKeys.EVAL).take(NUM_EVAL_EXAMPLES//1000)\n",
    "\n",
    "steps_per_epoch = NUM_TRAIN_EXAMPLES // (TRAIN_BATCH_SIZE * NUM_EVALS)\n",
    "\n",
    "history = model.fit(trainds, \n",
    "                    validation_data=evalds,\n",
    "                    epochs=NUM_EVALS, \n",
    "                    steps_per_epoch=steps_per_epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Visualize loss curve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot\n",
    "import matplotlib.pyplot as plt\n",
    "nrows = 1\n",
    "ncols = 2\n",
    "fig = plt.figure(figsize=(10, 5))\n",
    "\n",
    "for idx, key in enumerate(['loss', 'rmse']):\n",
    "    ax = fig.add_subplot(nrows, ncols, idx+1)\n",
    "    plt.plot(history.history[key])\n",
    "    plt.plot(history.history['val_{}'.format(key)])\n",
    "    plt.title('model {}'.format(key))\n",
    "    plt.ylabel(key)\n",
    "    plt.xlabel('epoch')\n",
    "    plt.legend(['train', 'validation'], loc='upper left');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import shutil, os, datetime\n",
    "OUTPUT_DIR = 'babyweight_trained'\n",
    "shutil.rmtree(OUTPUT_DIR, ignore_errors=True)\n",
    "EXPORT_PATH = os.path.join(OUTPUT_DIR, datetime.datetime.now().strftime('%Y%m%d%H%M%S'))\n",
    "tf.saved_model.save(model, EXPORT_PATH) # with default serving function\n",
    "print(\"Exported trained model to {}\".format(EXPORT_PATH))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls $EXPORT_PATH"
   ]
  },

  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Copyright 2020 Google Inc. Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
